"""
Tool agent implementation for handling tool-based operations
"""
from typing import Any, Dict, List, Optional
import json
import openai
from .base import BaseAgent
from ..core.message import Message, MessageType, ToolCallMessage, ToolResultMessage
from ..core.toolset import McpToolset

class ToolAgent(BaseAgent):
    """工具代理类，支持工具调用的智能体"""
    
    def __init__(self, 
                 agent_id: str,
                 api_key: str,
                 model: str = "gpt-3.5-turbo",
                 system_prompt: str = "你是一个具备工具调用能力的AI助手。",
                 duty: str = "协助用户完成任务。",
                 toolset: Optional[McpToolset] = None,
                 temperature: float = 0.7,
                 base_url: Optional[str] = None, 
                 max_num_context: int = 1000) -> None:
        """
        初始化工具代理
        Args:
            agent_id: 代理ID
            api_key: OpenAI API密钥
            model: 使用的模型名称
            toolset: 工具集实例
            temperature: 模型温度参数
            base_url: API基础URL，用于支持其他兼容OpenAI接口的模型服务
        """
        super().__init__(agent_id, max_num_context)
        self.client = openai.OpenAI(
            api_key=api_key,
            base_url=base_url if base_url else "https://api.openai.com/v1"
        )
        self.model = model
        self.system_prompt = system_prompt
        self.duty = duty
        self.toolset = toolset or McpToolset()
        self.temperature = temperature
        
    def _add_to_context(self, message: Dict) -> None:
        """将消息添加到会话上下文"""
        self.add_context(
            content=json.dumps(message)
        )
        
    def forward(self, message: Message) -> Message:
        """
        处理接收到的消息
        消息流程：
        1. 用户消息 -> Tool Agent
           - 如果需要调用工具：发送工具调用消息给自己
           - 如果不需要：直接返回响应给用户
        2. 工具调用消息 -> Tool Agent (self)
           - 执行工具调用
           - 将结果提供给模型
           - 返回最终响应给原始用户

        Args:
            message: 输入消息
        Returns:
            响应消息
        """
        if message.type == MessageType.TOOL_CALL:
            # 第2阶段：处理工具调用
            result = self._handle_tool_call(message)
            
            # 记录工具调用结果
            self._add_to_context({
                "role": "tool",
                "content": str(result.content["result"]) if result.content["success"] else str(result.content["error"]),
                "tool_call_id": message.metadata.get("tool_call_id")
            })
            # print("tool call id: ", message.metadata.get("tool_call_id"))
            
            # 让模型处理工具调用结果
            response = self._call_openai()
            
            # 记录模型的最终回复
            self._add_to_context({
                "role": "assistant",
                "content": response.choices[0].message.content
            })
            
            # 返回给原始用户
            return self.create_message(
                content=response.choices[0].message.content,
                msg_type=MessageType.TEXT,
                # receiver=message.metadata.get("original_sender")
            )
        else:
            # 第1阶段：处理用户消息
            # 记录用户消息
            self._add_to_context({
                "role": "user",
                "content": message.content
            })
            
            # 让模型判断是否需要使用工具
            response = self._call_openai()
            
            if response.choices[0].message.tool_calls:
                # 需要使用工具
                tool_call = response.choices[0].message.tool_calls[0]
                
                # 记录模型要使用工具的决定
                self._add_to_context(response.choices[0].message.model_dump())
                
                # 创建工具调用消息
                return ToolCallMessage(
                    tool_name=tool_call.function.name,
                    arguments=json.loads(tool_call.function.arguments),
                    sender=self.agent_id,
                    receiver=self.agent_id,
                    metadata={
                        "tool_call_id": tool_call.id,
                        "original_sender": message.sender
                    }
                )
            
            # 不需要使用工具
            # 记录模型的回复
            response_message = response.choices[0].message.model_dump()
            self._add_to_context(response_message)
            
            # 返回响应
            return self.create_message(
                content=response_message["content"],
                msg_type=MessageType.TEXT,
            )
        
    def _handle_tool_call(self, message: ToolCallMessage) -> ToolResultMessage:
        """处理工具调用请求"""
        tool_name = message.content["tool_name"]
        arguments = message.content["arguments"]
        
        try:
            # 执行工具调用
            result = self.toolset.tool_use(tool_name, **arguments)
            return ToolResultMessage(
                success=True,
                result=result,
                sender=self.agent_id,
                receiver=self.agent_id,
                metadata=message.metadata  # 保持元数据不变
            )
        except Exception as e:
            error_msg = str(e)
            print(f"工具调用出错 '{tool_name}': {error_msg}")
            return ToolResultMessage(
                success=False,
                result=None,
                error=error_msg,
                sender=self.agent_id,
                receiver=self.agent_id,
                metadata=message.metadata  # 保持元数据不变
            )
            
    def _call_openai(self) -> openai.types.chat.ChatCompletion:
        """调用OpenAI API"""
        # 构建消息历史
        messages = []
        
        # 添加系统提示
        messages.append({
            "role": "system",
            "content": self.system_prompt
        })
        
        # 添加会话历史
        for ctx in self.get_context():
            try:
                msg = json.loads(ctx["content"])
                messages.append(msg)
            except:
                continue
                
        # 获取工具定义
        tools = self.toolset.get_schemas()
        
        # 调用API
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                tools=tools,
                tool_choice="auto",  # 让模型自行决定是否使用工具
                temperature=self.temperature
            )
            return response
        except Exception as e:
            raise ValueError(f"Error calling OpenAI API: {str(e)}")
